#ifndef __MULTIARRAYS_H__
#define __MULTIARRAYS_H__

#include <vector>
#include <iostream>
#include <exception>
#include <cmath>
#include "Config.h"
#include "Vec.h"


/**
 * @brief A data structure that store multiple arrays
 * and each array coresponding to an grid in multi-grid 
 * method.
 * 
 * @tparam T 
 * @tparam N 
 */
template <typename T, int N>
class MultiArrays
{
private:
    /**
     * @brief store the data of multiple arrays
     * 
     */
    std::vector<T> _arrays;

    /**
     * @brief the number of levels or the number of arrays
     * 
     */
    int num_levels = 0;

    /**
     * @brief the number of data each array contains
     * 
     */
    std::vector<int> level_sizes; 

    /**
     * @brief the number of intervals that the finest grid have along each dimension
     * 
     */
    int _n;
public:

    MultiArrays() = default;

    /**
     * @brief Construct a new Multi Arrays object
     * 
     * @param n : the finest size of grid
     * @param num_levels : the maximum level of arrays
     */
    MultiArrays(int n, int num_levels);
    /**
     * @brief Construct a new Multi Arrays object
     * The difference between this and last constructor 
     * is that this constructor will create as many as
     * possible levels but make sure the coarsest grid
     * have grid_size at least 4.
     * @param n 
     */
    MultiArrays(int n);

    ~MultiArrays() = default;

    /**
     * @brief Get the begin iterator of each level.
     *          Level(num_levels) == end(), Level(0) == the begin of the finest grid
     * @param l 
     * @return std::vector<T>::iterator 
     */
    typename std::vector<T>::iterator Level(int l);
    typename std::vector<T>::const_iterator Level(int l) const;
    typename std::vector<T>::const_iterator end() const;

    /**
     * @brief Resize the multiarrays to new grids with 
     * given finest grid size.
     * 
     * @param new_n : the new finest grid size
     */
    void resize(int new_n);


    int getFinestSize() const { return _n; }

    int getLevelPointNum(int i) { return level_sizes[i]; }

    int getLevelNum() const { return num_levels; }

    /**
     * @brief Get the true idx of a grid points in the finest array.
     * 
     * @param idx : the idx of the point in the current array
     * @param l : the level of the array
     * @return int : the true idx in the finest grid
     */
    int getTrueIdx(int idx, int l);

    friend std::ostream &operator<<(std::ostream &out, const MultiArrays<T, N> &arrays)
    {
        for (int i = 0; i < arrays.num_levels; ++i)
        {
            out << "Level[" << i << "] :";
            for (auto it = arrays.Level(i); it != arrays.Level(i + 1); ++it)
            {
                out << *it << ",";
            }
            out << std::endl;
        }
        return out;
    }

};

/**
 * @brief Convert the idx in an 1d array to the coordinate in a grid with given grid size.
 * In the array, the grid points are sorted in lexicographical order w.r.t. their coordiantes.
 * 
 * @tparam N : dimension
 * @param idx 
 * @param grid_size : the number of intervals along each dimension
 * @return Vec<int, N> the coresponding coordinate
 */
template <int N>
Vec<int, N> idx2coord(int idx, int grid_size);

/**
 * @brief Convert a coordinate to the idx in the array. 
 * 
 * @tparam N 
 * @param coord 
 * @param grid_size 
 * @return int 
 */
template <int N>
int coord2idx(const Vec<int, N> &coord, int grid_size);

/**
 * @brief Check if a point with given index is 
 * on the boundary of current regular grid.
 *  
 * 
 * @tparam N 
 * @param idx 
 * @param grid_size 
 * @return true : on the boundary
 */
template <int N>
bool RegularBoundaryCheck(int idx, int grid_size);

template <typename T, int N>
MultiArrays<T, N>::MultiArrays(int n, int num_levels) : _n(n), _arrays(2 * std::pow(n + 1, N)), num_levels(num_levels), level_sizes(num_levels)
{
    int current_n = _n;
    for (int i = 0; i < num_levels; ++i, current_n /= 2)
    {
        level_sizes[i] = std::pow(current_n + 1, N);
    }
}

template <typename T, int N>
MultiArrays<T, N>::MultiArrays(int n) : _n(n), _arrays(2 * std::pow(n + 1, N))
{
    num_levels = 1;
    while (n > 4)
    {
        n /= 2;
        num_levels++;
    }
    level_sizes = std::vector<int>(num_levels);
    int current_n = _n;
    for (int i = 0; i < num_levels; ++i, current_n /= 2)
    {
        level_sizes[i] = std::pow(current_n + 1, N);
    }
}

template <typename T, int N>
void MultiArrays<T, N>::resize(int new_n)
{
    if (new_n != _n)
    {
        _n = new_n;
        num_levels = 1;
        while (new_n > 4)
        {
            new_n /= 2;
            num_levels++;
        }
        _arrays.resize(2 * std::pow(_n + 1, N));
        level_sizes.resize(num_levels);
    }
    level_sizes = std::vector<int>(num_levels);
    int current_n = _n;
    for (int i = 0; i < num_levels; ++i, current_n /= 2)
    {
        level_sizes[i] = std::pow(current_n + 1, N);
    }
}

template <typename T, int N>
typename std::vector<T>::iterator MultiArrays<T, N>::Level(int l)
{
    if (l > num_levels)
        throw std::out_of_range("The level doesnot exist!");
    int start_idx = 0;
    for (int i = 0; i < l; ++i)
    {
        start_idx += level_sizes[i];
    }
    return _arrays.begin() + start_idx;
}

template <typename T, int N>
typename std::vector<T>::const_iterator MultiArrays<T, N>::Level(int l) const
{
    if (l > num_levels)
        throw std::out_of_range("The level doesnot exist!");
    int start_idx = 0;
    for (int i = 0; i < l; ++i)
    {
        start_idx += level_sizes[i];
    }
    return _arrays.begin() + start_idx;
}

template <typename T, int N>
typename std::vector<T>::const_iterator MultiArrays<T, N>::end() const
{
    return this->Level(num_levels);
}

template <typename T, int N>
int MultiArrays<T, N>::getTrueIdx(int idx, int l)
{
    Vec<int, N> true_coord = idx2coord<N>(idx, _n >> l);
    for (int i = 0; i < l; ++i)
    {
        true_coord = true_coord * 2;
    }
    return coord2idx(true_coord, _n);
}

template <int N>
Vec<int, N> idx2coord(int idx, int grid_size)
{
    Vec<int, N> coord;
    for (int i = 0; i < N; ++i)
    {
        coord[i] = idx % (grid_size + 1);
        idx /= (grid_size + 1);
    }
    return std::move(coord);
}

template <int N>
bool RegularBoundaryCheck(int idx, int grid_size)
{
    int tmp = 1;
    for (int i = 0; i < N; ++i)
    {
        tmp = idx % (grid_size + 1);
        if (tmp == 0 || tmp == grid_size)
            return true;
        idx /= (grid_size + 1);
    }
    return false;
}

template <int N>
int coord2idx(const Vec<int, N> &coord, int grid_size)
{
    int idx = 0;
    for (int i = N - 1; i >= 0; --i)
    {
        idx = idx * (grid_size + 1) + coord[i];
    }
    return idx;
}
#endif // __MULTIARRAYS_H__